from manim import *
import numpy as np

class GraphGridVisualization(Scene):
    def construct(self):
        tille=Text("ABC-393-C-Make it sample")
        self.play(Write(tille))
        self.wait(2)
        self.clear()
        # 输入数据
        inputs = [
            (6, 10), (6, 2), (4, 1), (5, 1),
            (6, 6), (5, 3), (5, 1), (1, 4),
            (6, 4), (4, 2), (5, 6), 
        ]

        # 定义节点的数量和大小
        grid_size = len(set([item for sublist in inputs for item in sublist]))  # 获取唯一节点数量
        node_radius = 0.3
        area_side = int(np.sqrt(grid_size))
        if area_side ** 2 < grid_size:
            area_side += 1
        
        # 创建节点位置列表
        positions = []
        for i in range(11):
            x = (i % area_side) * node_radius * 4 - ((area_side - 1) / 2) * node_radius * 4
            y = -(i // area_side) * node_radius * 4 + ((area_side - 1) / 2) * node_radius * 4
            positions.append([x, y, 0])

        # 创建并放置节点
        nodes = {}
        for i in range(1, 11):
            circle = Circle(radius=node_radius, color=WHITE).move_to(positions[i-1])
            label = Text(str(i), font_size=24).next_to(circle, DOWN, buff=0.2)
            nodes[i] = VGroup(circle, label)

        s = {}  # 模拟图结构
        edge_counts = {}  # 记录每条边出现次数
        self.clear()  # 清除当前场景
        for i, (a, b) in enumerate(inputs):
            # 显示当前输入            
            current_input = Text("输入 {}: ({}, {})".format(i + 1, a, b), font_size=32).to_edge(UP)
            self.play(Write(current_input))

            if a > b: a, b = b, a  # 如果起点大于终点就交换
            key = (a, b)
            if key not in edge_counts: edge_counts[key] = 0
            edge_counts[key] += 1
            
            if edge_counts[key] > 1:  # 如果是重边
                line_color = RED  # 高亮颜色
            else:
                line_color = BLUE
                
            if a not in s: s[a] = set()
            s[a].add(b)  # 插入边
            
            
            
            # 动态创建图形表示
            self.add(*[nodes[k] for k in nodes])  # 添加所有节点
            for u in s.keys():
                for v in s[u]:
                    if u == v:  # 自边
                        radius = 0.5  # 圆圈的半径
                        loop_angle = 270  # 箭头所在的角度位置
                        
                        # 绘制自边的圆圈
                        arc = Arc(start_angle=PI/2, angle=-loop_angle*DEGREES, radius=radius,color=YELLOW)
                        arc.shift(nodes[u][0].get_center())
                        
                        self.add(arc)
                    else:
                        if (u, v) in edge_counts and edge_counts[(u, v)] > 1:
                            line_color = RED  # 重边高亮
                        else:
                            line_color = BLUE
                        if(nodes.get(u) and nodes.get(v)):
                            arrow = Arrow(nodes[u][0].get_center(), nodes[v][0].get_center(), color=line_color)
                            self.add(arrow)
            
            
            self.wait(2)
            self.play(FadeOut(current_input))

        # 计算并显示要删除的边数
        cnt = sum(len(edges) for edges in s.values())-1
        result_text = Text("要删除的边数：{}".format(len(inputs) - cnt), font_size=36).to_edge(UP)
        self.play(Write(result_text))
        self.wait(2)